import requests
import json
import logging
import time
from typing import Optional
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from config import settings  # 导入配置

logger = logging.getLogger(__name__)


class VolcanoLLM:
    def __init__(self, api_key: str = None, model_name: str = None):
        # 优先使用传入的参数，若无则使用配置中的默认值
        self.api_key = api_key or settings.volcano_api_key
        self.model_name = model_name or settings.volcano_model
        self.api_url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
        logger.debug(f"火山模型客户端初始化完成，使用模型: {self.model_name}")

    def generate(self, system_prompt: str, user_prompt: str, max_retries: int = 3) -> Optional[str]:
        """调用火山模型生成文本"""
        logger.debug(f"开始调用火山模型 {self.model_name}，系统提示词长度：{len(system_prompt)}，用户提示词长度：{len(user_prompt)}")

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            "temperature": 0.1,  # 低温度以获得更确定的输出
            "max_tokens": 4000  # 根据需求调整
        }

        # 创建会话并配置重试机制
        session = requests.Session()

        # 定义重试策略
        retry_strategy = Retry(
            total=max_retries,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["POST"]
        )

        # 创建HTTP适配器并挂载到会话
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("https://", adapter)
        session.mount("http://", adapter)

        # 禁用会话的环境变量代理信任
        session.trust_env = False

        for attempt in range(max_retries):
            try:
                logger.debug(f"尝试第 {attempt + 1} 次API调用...")

                # 发送POST请求
                response = session.post(
                    self.api_url,
                    headers=headers,
                    json=payload,
                    timeout=120,
                    proxies={"http": None, "https": None}
                )

                # 检查HTTP状态码
                if response.status_code == 404:
                    error_msg = "404错误：请求的资源未找到。请检查URL和模型名称是否正确。"
                    logger.error(error_msg)
                    return error_msg

                response.raise_for_status()

                # 解析响应JSON
                result = response.json()
                logger.debug(f"火山模型API返回：{json.dumps(result, ensure_ascii=False)}")

                # 提取模型生成的内容
                if "choices" in result and len(result["choices"]) > 0:
                    return result["choices"][0]["message"]["content"]
                else:
                    logger.error(f"火山模型响应格式异常：{result}")
                    return None

            except requests.exceptions.HTTPError as e:
                logger.error(f"HTTP错误 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    wait_time = 5 * (attempt + 1)
                    logger.debug(f"等待{wait_time}秒后重试...")
                    time.sleep(wait_time)
                else:
                    return f"处理失败: {str(e)}"
            except requests.exceptions.RequestException as e:
                logger.error(f"请求异常 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    wait_time = 5 * (attempt + 1)
                    logger.debug(f"等待{wait_time}秒后重试...")
                    time.sleep(wait_time)
                else:
                    return f"处理失败: {str(e)}"
            except KeyError as e:
                logger.error(f"解析响应时出错 (尝试 {attempt + 1}/{max_retries}): 响应中缺少预期的键 {e}")
                return f"处理失败: 无法解析API响应 - {str(e)}"
            except Exception as e:
                logger.error(f"未知错误 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    wait_time = 5 * (attempt + 1)
                    logger.debug(f"等待{wait_time}秒后重试...")
                    time.sleep(wait_time)
                else:
                    return f"处理失败: {str(e)}"

        return "处理失败: 超过最大重试次数"
